# -*- coding: UTF-8 -*-
import mlflow.pyfunc as pyfunc
import numpy as np
import base64
import cv2
import os
import sys
from ultralytics import YOLO
import logging
import torch
import time
import requests

logging.basicConfig(format='%(asctime)s.%(msecs)03d [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s',
                    datefmt='## %Y-%m-%d %H:%M:%S')
logging.getLogger().setLevel(logging.DEBUG)
logger = logging.getLogger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class model_predict(pyfunc.PythonModel):
    def load_context(self, context):
        curr_dir = os.path.dirname(os.path.abspath(__file__))
        sys.path.append(curr_dir)
        pdir = os.path.abspath(os.path.join(curr_dir, '../'))
        sys.path.insert(0, pdir)
        pdir2 = os.path.abspath(os.path.join(curr_dir, '../artifacts'))
        sys.path.insert(0, pdir2)
        self.artfact_path = pdir2
        weights_path = os.path.join(pdir2, 'best.pt')
        self.model = YOLO(weights_path)

        logger.info(f'---------预热开始...')
        for i in range(10):
            data = cv2.imdecode(np.fromfile(os.path.join(pdir2, '1.jpg'), dtype=np.uint8), cv2.IMREAD_COLOR)
            # data = cv2.imdecode(np.fromfile(context.artifacts["testpic_path"], dtype=np.uint8), cv2.IMREAD_COLOR)
            self.predict(context=context, inputs=data)
        logger.info(f'---------预热结束.')

    def predict(self, context, inputs):
        if isinstance(inputs, np.ndarray):
            image = inputs
            result = self.model(image)
        else:
            base64_str = inputs.iloc[0, 0]
            imgString = base64.b64decode(base64_str)
            nparr = np.fromstring(imgString, np.uint8)
            img_np = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            result = self.model(img_np)

        # 提取分割结果
        masks = result['masks']  # 分割掩码
        scores = result['scores']  # 置信度
        labels = result['labels']  # 标签
        boxes = result['boxes']  # 边界框

        nums = []
        for score, label, box, mask in zip(scores, labels, boxes, masks):
            # 将掩码转换为二值图像
            mask_image = (mask * 255).astype(np.uint8)
            _, contours, _ = cv2.findContours(mask_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # 提取轮廓点
            contour_points = []
            for contour in contours:
                contour_points.extend(contour.reshape(-1, 2).tolist())

            nums.append({
                'title': label,
                'type': 'polygon',
                'confidence': float(score),
                'shape': {
                    "points": contour_points
                }
            })

        return {"object": nums}
